package cn.zhxu.research;

import cn.zhxu.data.Array;
import cn.zhxu.data.Mapper;
import cn.zhxu.okhttps.HTTP;
import cn.zhxu.okhttps.HttpResult;
import cn.zhxu.okhttps.OkHttps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.text.DecimalFormat;
import java.time.Duration;
import java.util.List;
import java.util.stream.Stream;

public class Similarity {

    static final Logger log = LoggerFactory.getLogger(Similarity.class);

    static final int BACH_SIZE = 10;
    static final DecimalFormat SIMILARITY_FORMAT = new DecimalFormat("0.0000");

    static final HTTP http = OkHttps.getHttp()
            .newBuilder()
            .config(b -> {
                b.connectTimeout(Duration.ofSeconds(60));
                b.writeTimeout(Duration.ofSeconds(60));
                b.readTimeout(Duration.ofSeconds(300));
            })
            .build();

    static List<float[]> embeddings(String model, List<String> texts) {
        var res = http.sync("http://localhost:9997/v1/embeddings")
                .bodyType(OkHttps.JSON)
                .addBodyPara("model", model)
                .addBodyPara("input", texts)
                .post();
        if (res.getState() != HttpResult.State.RESPONSED) {
            log.info(res.toString());
            throw new RuntimeException("Embedding No Response!");
        }
        Mapper body = res.getBody().toMapper();
        int status = res.getStatus();
        if (status != 200) {
            log.info(res.toString());
            log.info(body.toString());
            throw new RuntimeException("Embedding Response Status error: " + status);
        }
        return body.getArray("data").stream()
                .map(data -> data.toMapper().getArray("embedding"))
                .map(Array::toFloatArray)
                .toList();
    }


    public static void main(String[] args) throws IOException {
        for (String model : Models.models) {
            process(model);
        }
    }

    static void process(String model) throws IOException {
        long t0 = System.currentTimeMillis();
        var dataItems = DataItem.readDataItems("E:/SpaceMine/research/data/dataset.txt");
        var totalCnt = dataItems.size();
        int processedCnt = 0;
        try (BufferedWriter writer = new BufferedWriter(new FileWriter("E:/SpaceMine/research/data/models2/" + model + ".txt"))) {
            while (processedCnt < totalCnt) {
                int count = Math.min(BACH_SIZE, totalCnt - processedCnt);
                log.info("Embedding {} ~ {}", processedCnt, processedCnt + count);
                var items = dataItems.subList(processedCnt, processedCnt + count);
                doEmbedding(model, writer, items);
                processedCnt += count;
            }
        }
        long t1 = System.currentTimeMillis();
        log.info("Embedding Time: {} min - {}", (t1 - t0) / 1000 / 60, model);
    }

    static void doEmbedding(String model, BufferedWriter writer, List<DataItem> items) throws IOException {
        var texts = items.stream()
                .flatMap(it -> Stream.of(it.getTxt1(), it.getTxt2()))
                .toList();
        var embeddings = embeddings(model, texts);
        if (embeddings.size() != texts.size()) {
            log.info("Embedding 数量不对 {} != {}", embeddings.size(), texts.size());
            return;
        }
        for (int i = 0; i < items.size(); i++) {
            var item = items.get(i);
            var v1 = embeddings.get(2 * i);
            var v2 = embeddings.get(2 * i + 1);
            var similarity = cosineSimilarity(v1, v2);
            var distance = euclideanDistance(v1, v2);
            writer.write(item.getScore() + "##" + similarity + "##" + distance);
            writer.newLine();
        }
    }

    static double innerProduct(float[] v1, float[] v2) {
        if (v1.length != v2.length) {
            throw new IllegalStateException("Vectors must be of the same length");
        }
        double sum = 0;
        for (int i = 0; i < v1.length; i++) {
            sum += v1[i] * v2[i];
        }
        return sum;
    }

    static double magnitude(float[] v) {
        double sum = 0;
        for (float value : v) {
            sum += value * value;
        }
        return Math.sqrt(sum);
    }
    static double cosineSimilarity(float[] v1, float[] v2) {
        if (v1.length != v2.length) {
            throw new Error("Vectors must be of the same length");
        }
        var magA = magnitude(v1);
        var magB = magnitude(v2);
        if (magA == 0 || magB == 0) {
            return 0;
        }
        return innerProduct(v1, v2) / (magA * magB);
    }

    static double euclideanDistance(float[] v1, float[] v2) {
        if (v1.length != v2.length) {
            throw new Error("Vectors must be of the same length");
        }
        double sum = 0;
        for (int i = 0; i < v1.length; i++) {
            sum += Math.pow(v1[i] - v2[i], 2);
        }
        return Math.sqrt(sum);
    }

}
